from config import BASE_DIR
from get_dy_hot.get_dy_hot import get_dy_hot, login
from langchain_core.tools import tool
from langchain_community.llms import Tongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv
import os

# 加载环境变量
load_dotenv()

# 初始化通义千问模型
api_key = os.getenv("DASHSCOPE_API_KEY")
llm = Tongyi(
    model="qwen-turbo",
    api_key=api_key,
    temperature=0.7,
    max_tokens=2048
)

@tool
def get_dy_hot_tool():
    """获取抖音热搜和实时上升热搜信息"""
    try:
        today_hot, upping_hot = get_dy_hot()
        return f"今日热搜: {today_hot}\n实时上升热搜: {upping_hot}"
    except Exception as e:
        return f"获取热搜信息失败: {str(e)}"

@tool
def login_tool():
    """使用DrissionPage打开浏览器登录抖音页面"""
    try:
        login()
        return "抖音登录页面已成功打开"
    except Exception as e:
        return f"打开登录页面失败: {str(e)}"

# 工具列表
tools = [login_tool, get_dy_hot_tool]

def analyze_and_call_tool(user_input):
    """分析用户输入并调用相应工具"""
    # 使用LLM分析用户意图
    analysis_prompt = f"""
用户输入: {user_input}

请分析用户的意图，并选择合适的操作：
1. 如果用户想要打开抖音登录页面，回复"LOGIN"
2. 如果用户想要获取抖音热搜信息，回复"HOTSEARCH"
3. 如果是其他请求，回复"OTHER"

只回复操作类型，不要其他内容。
"""
    
    try:
        analysis_result = llm.invoke(analysis_prompt)
        print(f"意图分析结果: {analysis_result}")
        
        if "LOGIN" in analysis_result.upper():
            print("检测到登录请求，调用login_tool...")
            result = login_tool.invoke({})
            return result
        elif "HOTSEARCH" in analysis_result.upper():
            print("检测到热搜查询请求，调用get_dy_hot_tool...")
            result = get_dy_hot_tool.invoke({})
            return result
        else:
            return "抱歉，我无法理解您的请求。我可以帮您打开抖音登录页面或获取热搜信息。"
            
    except Exception as e:
        return f"分析用户意图时出错: {str(e)}"

def main():
    """主函数 - 测试工具调用"""
    print("=== LangChain工具调用测试 ===")
    
    # 第一次对话：测试登录工具
    print("\n--- 第一次对话：测试登录功能 ---")
    user_input1 = "请帮我打开抖音登录页面"
    print(f"用户输入: {user_input1}")
    
    result1 = analyze_and_call_tool(user_input1)
    print(f"执行结果: {result1}")
    
    print("\n" + "="*50)
    
    # 第二次对话：测试热搜工具
    print("\n--- 第二次对话：测试热搜功能 ---")
    user_input2 = "今日的抖音热搜信息和实时上升热搜都有哪些"
    print(f"用户输入: {user_input2}")
    
    result2 = analyze_and_call_tool(user_input2)
    print(f"\n=== 热搜信息 ===")
    print(result2)
    
    print("\n=== 测试完成 ===")

if __name__ == "__main__":
    main()

